#这里包含densenet网络的变体
from __future__ import absolute_import#兼容python3.x和2.x，加入绝对引入
from __future__ import division#导入python未来支持的语言特征division(精确除法)
from __future__ import print_function #兼容python2.x和3.x的print用法

import tensorflow as tf#导入tensorflow模块

slim = tf.contrib.slim#非官方，在未加入核心库之前接口可能会被修改，或还需要测试


def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)#从截断的正态分布中输出随机值，生成的值
 #服从具有指定平均值和标准偏差的正态分布，如果生成的值大于平均值2个标准偏差的值则丢弃重新选择。


def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):#H复杂函数
    #num_outputs为输出通道数,
    current = slim.batch_norm(current, scope=scope + '_bn')#批归一化,在激活函数之前用
    current = tf.nn.relu(current)#激活函数
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')#卷积
    current = slim.dropout(current, scope=scope + '_dropout')#减小过拟合，在分类器之前用
    return current


def block(net, layers, growth, scope='block'):#DenseBlock
    for idx in range(layers):
        #瓶颈层
        #1x1的卷积核的使用可以有效降低feature_map的数量，且将特征图的数量统一压缩为4*增长率
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1],scope=scope + '_conv1x1' + str(idx))
        
        #使用Ｈ函数,进行非线性变换，特征图转换为增长率的值的个数
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3],scope=scope + '_conv3x3' + str(idx))
        #稠密链接,每次增加growth个特征图在channel的维度上
        net = tf.concat(axis=3, values=[net, tmp])
    return net#返回net+tmp的值的数量和


def transition(net, num_outputs, scope='transition'):#转换层,1个卷积层和1个2*2的平均池化层
    net = bn_act_conv_drp(net, num_outputs, [1, 1], scope=scope + '_conv1x1')
    net = slim.avg_pool2d(net, [2, 2], stride=2, scope=scope + '_avgpool')
    return net

def densenet(images, num_classes=1001, is_training=False,dropout_keep_prob=0.8,scope='densenet'):#实验有3个dense block，每个block 有相同的层
   
    growth = 24#增长率指定
    compression_rate = 0.5#压缩因子指定，因为复现，按照论文中来，压缩因子小于1，加上上面的BN-ReLU-Conv(1x1)-BN-ReLU-Conv(3x3)，现在此网络为DenseNet-BC

    def reduce_dim(input_feature):#减少特征图的数量
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(bn_drp_scope(is_training=is_training,keep_prob=dropout_keep_prob)) as ssc:

            #因为是DenseNet-BC,所以输出通道数按照２倍的增长率；首先用7*7卷积核，3*3的最大池化
            net = slim.conv2d(images, 2*growth, [7, 7], stride=2, padding='same', scope='init_conv')
            end_points['init_conv'] = net
            net = slim.max_pool2d(net, [3, 3], stride=2, scope='init_conv')
            end_points['init_conv'] = net

            #增长率小则可能信息提取不出来，太大则导致特征图增加太快，运行速度变慢
            net = block(net, 10, growth, scope='block1')
            end_points['block1'] = net
            
            #1转换层
            net = transition(net, reduce_dim(net), scope='transition1')
            end_points['transition1'] = net
            
            #第二个dense block
            net = block(net, 10, growth, scope='block2')
            end_points['block2'] = net
            
            #2转换层
            net = transition(net, reduce_dim(net), scope='transition2')
            end_points['transition2'] = net
            
            #第三个dense block
            net = block(net, 10, growth, scope='block3')
            end_points['block3'] = net
            
            #3转换层
            net = transition(net, reduce_dim(net), scope='transition3')
            end_points['transition3'] = net

            #加这个AuxLogits,会加速整个训练的收敛，也可防止过拟合。
            aux = bn_act_conv_drp(net, reduce_dim(net),[3,3], scope = "AuxLogits")
            aux = slim.conv_drop(net, reduce_dim(net), [3,3], activation_fn = tf.nn.relu, scope = "AuxLogits_conv1")
            aux = slim.avg_pool2d(aux, aux.shape[1:3])
            aux = slim.flatten(aux)
            aux_logits = slim.fully_connected(aux, num_classes, activation_fn = None, scope = "Aux_Logits" )
            end_points["AuxLogits"] = aux_logits

            #第四个dense block
            net = block(net, 10, growth, scope='block4')
            end_points['block4'] = net
            
            #获取此时卷积核的大小
            kernel_size = net.get_shape()[1:3]
            
            #GAP
            net = slim.avg_pool2d(net, kernel_size, padding='VALID', scope='global_pool')
            end_points['global_pool'] = net
            
            #Flatten
            net = slim.flatten(net, scope='flatten')
            end_points['flatten'] = net
            
            #fully_connected,logits
            logits = slim.fully_connected(net, num_classes,
                                  biases_initializer=tf.zeros_initializer(),
                                  weights_initializer=trunc_normal(0.1),
                                  scope='logits')
            end_points['Logits'] = logits
            #softmax
            end_points['predictions'] = slim.softmax(logits, scope='predictions')

    return logits, end_points


def bn_drp_scope(is_training=True, keep_prob=0.8):
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope([slim.batch_norm],scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope([slim.dropout], is_training=is_training, keep_prob=keep_prob) as bsc:
            return bsc


def densenet_arg_scope(weight_decay=0.004):
    with slim.arg_scope(
        [slim.conv2d],
        weights_initializer=tf.contrib.layers.variance_scaling_initializer(
            factor=2.0, mode='FAN_IN', uniform=False),
        activation_fn=None, biases_initializer=None, padding='same',
            stride=1) as sc:
        return sc


densenet.default_image_size = 224